import sys
import numpy as np
from PySide6.QtWidgets import QApplication, QMainWindow, QVBoxLayout, QWidget, QLabel, QGraphicsView
from PySide6.QtCharts import QChart, QChartView, QLineSeries, QValueAxis
from PySide6.QtGui import QPen, QColor
from PySide6.QtCore import Qt, QPointF

# 模拟数据
np.random.seed(0)
num_points = 50 # 增加数据点以便图表更明显
well_id = 'A1'
depths = np.linspace(1000, 2000, num_points)
rock_desc = ['砂岩', '泥岩', '灰岩']
analyst = '张三'
date = '2024-06-01'
types = ['a', 'b', 'c']
data = []

for i in range(num_points):
    row = {
        '井号': well_id,
        '深度': depths[i],
        '岩性简述': np.random.choice(rock_desc),
        '分析人': analyst,
        '分析日期': date,
        '类型': np.random.choice(types),
        'Mg': np.random.uniform(0, 15), # 调整范围使曲线变化更明显
        'Al': np.random.uniform(5, 20),
        'Si': np.random.uniform(10, 25),
        'P': np.random.uniform(0, 5),
        'S': np.random.uniform(0, 5),
        'Cl': np.random.uniform(0, 5),
        'K': np.random.uniform(0, 10),
        'Ca': np.random.uniform(5, 15),
        'Ti': np.random.uniform(0, 3),
        'V': np.random.uniform(0, 2),
        'Mn': np.random.uniform(0, 8),
        'Fe': np.random.uniform(10, 30),
    }
    data.append(row)

class ChartWidget(QMainWindow):
    def __init__(self, data):
        super().__init__()

        self.data = data
        self.setWindowTitle("井数据可视化 - Qt Charts 地质导向图")
        self.setGeometry(100, 100, 1000, 800)

        central_widget = QWidget()
        self.setCentralWidget(central_widget)

        layout = QVBoxLayout(central_widget)

        self.chart = QChart()
        self.chart.setTitle("地质导向交互折线图")
        self.chart.setAnimationOptions(QChart.SeriesAnimations)

        # 创建坐标轴
        self.axisX = QValueAxis()
        self.axisX.setTitleText("深度")
        self.axisX.setLabelFormat("%.1f")
        self.axisX.setTickCount(10) # 增加刻度线数量
        min_depth, max_depth = min([d['深度'] for d in data]), max([d['深度'] for d in data])
        self.axisX.setRange(min_depth, max_depth)

        self.axisY = QValueAxis()
        self.axisY.setTitleText("元素含量")
        self.axisY.setLabelFormat("%.1f")
        self.axisY.setTickCount(10) # 增加刻度线数量
        # 根据模拟数据调整纵坐标范围，考虑所有选定元素的范围
        elements_to_plot = ['Mg', 'Al', 'Si'] # 默认绘制的元素
        all_values = []
        for elem in elements_to_plot:
             all_values.extend([d[elem] for d in data])
        min_y, max_y = min(all_values) if all_values else 0, max(all_values) if all_values else 1
        self.axisY.setRange(min_y - (max_y-min_y)*0.1 , max_y + (max_y-min_y)*0.1)

        self.chart.addAxis(self.axisX, Qt.AlignBottom)
        self.chart.addAxis(self.axisY, Qt.AlignLeft)

        # 添加数据系列 (折线)
        self.series = {}
        colors = { 'Mg': QColor(255, 0, 0), 'Al': QColor(0, 150, 255), 'Si': QColor(0, 200, 0) } # 定义a/b/c对应的元素和颜色
        types_map = { 'a': 'Mg', 'b': 'Al', 'c': 'Si' } # 将类型映射到元素

        for chart_type, element in types_map.items():
             series = QLineSeries()
             series.setName(f'{chart_type} ({element})')
             pen = QPen(colors[element])
             pen.setWidth(2)
             series.setPen(pen)
             # 添加数据点 QPointF(x, y) -> QPointF(深度, 元素含量)
             for row in self.data:
                 series.append(QPointF(row['深度'], row[element]))

             self.chart.addSeries(series)
             series.attachAxis(self.axisX)
             series.attachAxis(self.axisY)
             self.series[chart_type] = series

             # 启用点点击信号，用于触碰显示节点数据
             series.clicked.connect(self.handlePointClicked) # 使用clicked信号作为触碰事件
             # Note: Qt Charts 没有内置的 hover 事件到具体点，clicked 是一个相对简单的方式演示触碰交互

        # 创建 ChartView
        self.chart_view = QChartView(self.chart)
        self.chart_view.setRenderHint(QPainter.Antialiasing)
        
        # 设置交互模式：拖动平移
        # self.chart_view.setRubberBand(QChartView.RectangleRubberBand) # 默认矩形缩放
        # setDragMode(QGraphicsView.ScrollHandDrag) 是 QGraphicsView 的方法，QChartView 继承自它
        self.chart_view.setDragMode(QGraphicsView.ScrollHandDrag) # 启用拖动手势进行平移

        layout.addWidget(self.chart_view)

        # 用于显示节点数据的 QLabel
        self.data_label = QLabel("悬停/点击节点显示数据")
        self.data_label.setAlignment(Qt.AlignCenter)
        layout.addWidget(self.data_label)

    # 处理点点击事件，显示数据
    def handlePointClicked(self, point):
        # 找到是哪个 series 的点被点击了
        clicked_series = self.sender() # sender() 返回发送信号的对象 (即 QLineSeries)
        if isinstance(clicked_series, QLineSeries):
             # 在 series 中找到被点击的点 (基于坐标)
             # 这是一个简化的查找，对于大量数据可能需要优化
             # 更好的方式是连接 QLineSeries 的 pointsPressed 或 pointsHovered 信号，但 PySide6 的 QLineSeries 没有直接提供简单的 point value 信号
             data_text = f"深度: {point.x():.1f}, 含量: {point.y():.2f}"

             # 更详细的数据查找（需要遍历原始数据）
             matching_rows = [row for row in self.data if abs(row['深度'] - point.x()) < 0.1 and abs(row[clicked_series.name().split(' (')[1][:-1]] - point.y()) < 0.1]
             if matching_rows:
                 row = matching_rows[0]
                 data_text = f"井号: {row['井号']}\n"
                 data_text += f"深度: {row['深度']:.1f}\n"
                 data_text += f"岩性: {row['岩性简述']}\n"
                 data_text += f"类型: {row['类型']}\n"
                 # 根据 series name 获取元素，然后显示该元素的值
                 element = clicked_series.name().split(' (')[1][:-1] # 从 "a (Mg)" 提取 "Mg"
                 data_text += f"{element}: {row[element]:.2f}\n"
                 data_text += f"分析人: {row['分析人']}\n"
                 data_text += f"日期: {row['分析日期']}"


             self.data_label.setText(data_text)


def main():
    app = QApplication(sys.argv)
    w = ChartWidget(data)
    w.show()
    sys.exit(app.exec())

if __name__ == '__main__':
    main() 